import os
import json
import numpy as np
from PIL import Image

def convert_labelme_to_yolo(json_file, output_dir):
    try:
        # 使用 UTF-8 编码读取 JSON 文件
        with open(json_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
    except UnicodeDecodeError:
        # 如果 UTF-8 编码读取失败，尝试使用 GBK 编码
        with open(json_file, 'r', encoding='gbk') as f:
            data = json.load(f)

    image_path = data['imagePath']
    image = Image.open(os.path.join(os.path.dirname(json_file), image_path))
    width, height = image.size

    shapes = data['shapes']
    txt_filename = os.path.splitext(os.path.basename(json_file))[0] + '.txt'
    txt_path = os.path.join(output_dir, txt_filename)

    with open(txt_path, 'w') as f:
        for shape in shapes:
            points = np.array(shape['points'])
            x_min = np.min(points[:, 0])
            y_min = np.min(points[:, 1])
            x_max = np.max(points[:, 0])
            y_max = np.max(points[:, 1])
            x_center = (x_min + x_max) / 2 / width
            y_center = (y_min + y_max) / 2 / height
            bbox_width = (x_max - x_min) / width
            bbox_height = (y_max - y_min) / height
            f.write(f"0 {x_center} {y_center} {bbox_width} {bbox_height}\n")

def main():
    input_dir = r"C:\Users\aqzz1\Desktop\对照组补json"
    output_dir = r"C:\Users\aqzz1\Desktop\对照组补txt"

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for filename in os.listdir(input_dir):
        if filename.endswith(".json"):
            json_path = os.path.join(input_dir, filename)
            convert_labelme_to_yolo(json_path, output_dir)
            print(f"已成功转换 '{filename}'")

if __name__ == "__main__":
    main()